#!/bin/bash

# ABACUS势能面扫描测试脚本
# ABACUS路径: /usr/local/bin/abacus

echo "=== ABACUS势能面扫描测试脚本 ==="
echo "ABACUS路径: /usr/local/bin/abacus"

# 设置
ABACUS_PATH="/usr/local/bin/abacus"
CORE=2  # 测试用2个核心
WORK_DIR="test_abacus_scan"
PSEUDO_DIR="./potential"  # 假设赝势文件在当前目录的potential文件夹中

# 检查ABACUS是否存在
echo "=== 环境检查 ==="
if [ ! -f "$ABACUS_PATH" ]; then
    echo "❌ 错误: ABACUS不存在于 $ABACUS_PATH"
    echo "请确认ABACUS已正确安装"
    exit 1
else
    echo "✓ ABACUS找到: $ABACUS_PATH"
fi

# 检查MPI
if command -v mpirun &> /dev/null; then
    echo "✓ MPI找到"
else
    echo "❌ 错误: mpirun未找到"
    exit 1
fi

# 创建工作目录
echo "=== 创建工作目录 ==="
rm -rf $WORK_DIR
mkdir -p $WORK_DIR
cd $WORK_DIR
echo "工作目录: $(pwd)"

# 创建INPUT文件
echo "=== 创建输入文件 ==="
cat > INPUT << 'EOF'
#INPUT_PARAMETERS
pseudo_dir      ./potential
orbital_dir     ./potential

#Parameters (1.General)
calculation     scf

#Parameters (2.Iteration)
ecutwfc         60
scf_nmax        30
dr2             1e-4

#Parameters (3.Basis)
basis_type      lcao
kspacing        0.15 0.15 0.2

#Parameters (4.Smearing)
smearing_method gauss
smearing_sigma  0.01

#Parameters (5.Mixing)
mixing_type     pulay
mixing_beta     0.7

cal_stress      1
cal_force       1
EOF
echo "✓ INPUT文件创建完成"

# 创建赝势目录（如果不存在）
mkdir -p potential

# 检查赝势文件
echo "=== 检查赝势文件 ==="
if [ -d "../potential" ]; then
    echo "复制赝势文件..."
    cp ../potential/*.upf potential/ 2>/dev/null || echo "注意: 未找到upf文件"
    cp ../potential/*.orb potential/ 2>/dev/null || echo "注意: 未找到orb文件"
fi

echo "赝势文件: $(ls potential/*.upf 2>/dev/null | wc -l) 个upf文件"
echo "轨道文件: $(ls potential/*.orb 2>/dev/null | wc -l) 个orb文件"

# 测试点坐标
TEST_POINTS=(
    "0.1 0.1 0.09"
    "0.2 0.2 0.09" 
    "0.3 0.1 0.09"
)

# 运行测试计算
echo "=== 开始测试计算 ==="
total_tests=${#TEST_POINTS[@]}
success_count=0

for ((i=0; i<${#TEST_POINTS[@]}; i++)); do
    coords=(${TEST_POINTS[$i]})
    x=${coords[0]}
    y=${coords[1]}
    z=${coords[2]}
    
    dir_name="test_${i}_x${x}_y${y}_z${z}"
    mkdir -p $dir_name
    cd $dir_name
    
    echo "--- 测试点 $((i+1))/$total_tests: ($x, $y, $z) ---"
    
    # 复制INPUT文件
    cp ../INPUT ./
    cp -r ../potential ./
    
    # 创建STRU文件 - 使用简化的小系统测试
    cat > STRU << EOF
ATOMIC_SPECIES
C   12.011  C_ONCV_PBE-1.0.upf
Al   26.982  Al_ONCV_PBE-1.0.upf

NUMERICAL_ORBITAL
C_gga_7au_100Ry_2s2p1d.orb
Al_gga_7au_100Ry_4s4p1d.orb

LATTICE_CONSTANT
1.889726

LATTICE_VECTORS
6.340000000000  0.000000000000  0.000000000000
-3.17000000000  5.490601059993  0.000000000000
0.000000000000  0.000000000000  15.00000000000

ATOMIC_POSITIONS
Direct

C 
0.000
6
0.259660000000  0.129830000000  0.000000000000  1  1  1
0.870170000000  0.129830000000  0.000000000000  1  1  1  
0.870170000000  0.740340000000  0.000000000000  1  1  1
0.740340000000  0.870170000000  0.000000000000  1  1  1
0.129830000000  0.870170000000  0.000000000000  1  1  1
0.129830000000  0.259660000000  0.000000000000  1  1  1

Al
0.000
1
$x  $y  $z  1  1  1
EOF
    
    echo "运行ABACUS计算..."
    start_time=$(date +%s)
    
    # 运行ABACUS计算
    export OMP_NUM_THREADS=1
    mpirun -n $CORE $ABACUS_PATH > out.log 2> err.log
    exit_code=$?
    
    end_time=$(date +%s)
    duration=$((end_time - start_time))
    
    # 检查计算结果
    if [ $exit_code -eq 0 ] && [ -f "OUT.ABACUS/running_scf.log" ]; then
        if grep -q "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log; then
            energy=$(grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log | tail -1 | awk '{print $2}')
            echo "✅ 测试成功! 能量: $energy Ry, 耗时: ${duration}秒"
            ((success_count++))
        else
            echo "❌ 计算未收敛"
            echo "最后10行输出:"
            tail -10 out.log
        fi
    else
        echo "❌ 计算失败, 退出码: $exit_code"
        echo "错误日志:"
        cat err.log
    fi
    
    cd ..
    echo ""
done

# 测试结果汇总
echo "=== 测试结果汇总 ==="
echo "总测试点: $total_tests"
echo "成功数量: $success_count"
echo "成功率: $((success_count * 100 / total_tests))%"

if [ $success_count -eq $total_tests ]; then
    echo "🎉 所有测试通过! 可以运行完整扫描脚本"
    
    # 生成简化的后续扫描建议
    echo ""
    echo "=== 后续操作建议 ==="
    echo "1. 确认INPUT参数设置正确"
    echo "2. 检查赝势文件完整性" 
    echo "3. 运行完整势能面扫描脚本"
    echo "4. 监控计算资源和进度"
    
else
    echo "⚠️ 部分测试失败，请检查:"
    echo "   - 赝势文件是否完整"
    echo "   - INPUT参数是否合适"
    echo "   - 系统资源是否充足"
fi

# 清理（可选）
echo ""
read -p "是否清理测试文件? (y/n): " cleanup
if [ "$cleanup" = "y" ] || [ "$cleanup" = "Y" ]; then
    cd ..
    rm -rf $WORK_DIR
    echo "测试文件已清理"
else
    echo "测试文件保留在: $WORK_DIR"
fi

echo "=== 测试脚本执行完成 ==="